import cv2
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path

class HandEyeCalibration:
    def __init__(self):
        self.pixel_points = []  # 像素坐标
        self.world_points = []  # 世界坐标（电机坐标）
        self.H = None  # 单应性矩阵
        self.rotation_matrix = None  # 旋转矩阵
        self.translation_vector = None  # 平移向量
        self.scale_factor = None  # 尺度因子
        
    def load_calibration_data_from_file(self, file_path):
        """从txt文件加载标定数据"""
        file_path = Path(file_path)
        if not file_path.exists():
            raise FileNotFoundError(f"标定文件不存在: {file_path}")
        
        print(f"正在读取标定文件: {file_path}")
        
        with open(file_path, 'r', encoding='utf-8') as f:
            lines = f.readlines()
        
        for line in lines:
            line = line.strip()
            if line.startswith('#') or line.startswith('pixel_x') or not line:
                continue  # 跳过注释行、标题行和空行
            
            parts = line.split(',')
            if len(parts) == 4:
                try:
                    pixel_x = float(parts[0])
                    pixel_y = float(parts[1])
                    world_x = float(parts[2])
                    world_y = float(parts[3])
                    
                    # 由于y方向相反，对世界坐标的y取反
                    self.pixel_points.append([pixel_x, pixel_y])
                    self.world_points.append([world_x, -world_y])  # y方向取反
                except ValueError as e:
                    print(f"解析行时出错: {line} - {e}")
        
        if not self.pixel_points:
            raise ValueError("未找到有效的标定数据")
        
        self.pixel_points = np.array(self.pixel_points, dtype=np.float32)
        self.world_points = np.array(self.world_points, dtype=np.float32)
        
        print(f"成功加载 {len(self.pixel_points)} 个标定点")
        print("像素坐标:")
        for i, (px, py) in enumerate(self.pixel_points):
            print(f"  点{i+1}: ({px:.1f}, {py:.1f})")
        
        print("世界坐标（电机坐标，y已取反）:")
        for i, (wx, wy) in enumerate(self.world_points):
            print(f"  点{i+1}: ({wx:.1f}, {wy:.1f})")
    
    def estimate_homography(self):
        """估计单应性矩阵"""
        # 使用RANSAC方法估计单应性矩阵
        self.H, mask = cv2.findHomography(
            self.pixel_points, self.world_points, cv2.RANSAC, 5.0
        )
        
        print("\n=== 单应性矩阵估计 ===")
        print(f"单应性矩阵 H:\n{self.H}")
        
        # 计算重投影误差
        self.calculate_reprojection_error()
        
        return self.H
    
    def calculate_scale_factor(self):
        """计算像素到世界坐标的尺度因子"""
        if len(self.pixel_points) < 2:
            return None
        
        # 计算像素距离
        pixel_distances = []
        world_distances = []
        
        for i in range(len(self.pixel_points)):
            for j in range(i + 1, len(self.pixel_points)):
                pixel_dist = np.linalg.norm(self.pixel_points[i] - self.pixel_points[j])
                world_dist = np.linalg.norm(self.world_points[i] - self.world_points[j])
                
                if world_dist > 0 and pixel_dist > 0:  # 避免除零
                    pixel_distances.append(pixel_dist)
                    world_distances.append(world_dist)
        
        # 计算平均尺度因子
        if world_distances:
            scale_factors = [wd / pd for pd, wd in zip(pixel_distances, world_distances)]
            self.scale_factor = np.mean(scale_factors)
            
            print(f"\n=== 尺度因子计算 ===")
            print(f"像素到世界的尺度因子: {self.scale_factor:.6f}")
            print(f"世界到像素的尺度因子: {1/self.scale_factor:.6f}")
        else:
            print("无法计算尺度因子：没有有效的距离对")
        
        return self.scale_factor
    
    def calculate_reprojection_error(self):
        """计算重投影误差"""
        if self.H is None:
            return
        
        # 使用单应性矩阵将像素坐标转换到世界坐标
        transformed_points = cv2.perspectiveTransform(
            self.pixel_points.reshape(-1, 1, 2), self.H
        ).reshape(-1, 2)
        
        # 计算误差
        errors = np.linalg.norm(transformed_points - self.world_points, axis=1)
        mean_error = np.mean(errors)
        max_error = np.max(errors)
        rmse_error = np.sqrt(np.mean(errors**2))
        
        print(f"\n=== 重投影误差分析 ===")
        print(f"平均重投影误差: {mean_error:.6f} 单位")
        print(f"最大重投影误差: {max_error:.6f} 单位")
        print(f"RMSE误差: {rmse_error:.6f} 单位")
        print("各点误差:")
        for i, error in enumerate(errors):
            print(f"  点{i+1}: {error:.6f}")
        
        return mean_error, max_error, rmse_error
    
    def pixel_to_world(self, pixel_x, pixel_y):
        """将像素坐标转换为世界坐标"""
        if self.H is None:
            raise ValueError("请先进行手眼标定")
        
        pixel_point = np.array([[pixel_x, pixel_y]], dtype=np.float32)
        world_point = cv2.perspectiveTransform(
            pixel_point.reshape(-1, 1, 2), self.H
        ).reshape(-1, 2)[0]
        
        return world_point[0], world_point[1]
    
    def world_to_pixel(self, world_x, world_y):
        """将世界坐标转换为像素坐标"""
        if self.H is None:
            raise ValueError("请先进行手眼标定")
        
        # 对y坐标取反（因为电机坐标系y方向相反）
        world_point = np.array([[world_x, -world_y]], dtype=np.float32)
        pixel_point = cv2.perspectiveTransform(
            world_point.reshape(-1, 1, 2), np.linalg.inv(self.H)
        ).reshape(-1, 2)[0]
        
        return pixel_point[0], pixel_point[1]
    
    def visualize_calibration(self):
        """可视化标定结果"""
        if self.H is None:
            print("请先进行标定")
            return
        
        # 使用单应性矩阵转换所有点
        transformed_points = cv2.perspectiveTransform(
            self.pixel_points.reshape(-1, 1, 2), self.H
        ).reshape(-1, 2)
        
        # 创建图表
        plt.figure(figsize=(15, 6))
        
        # 子图1：原始数据
        plt.subplot(1, 2, 1)
        plt.scatter(self.pixel_points[:, 0], self.pixel_points[:, 1], 
                   c='red', label='像素坐标', s=50)
        for i, (x, y) in enumerate(self.pixel_points):
            plt.annotate(f'{i+1}', (x, y), xytext=(5, 5), textcoords='offset points')
        plt.xlabel('像素 X')
        plt.ylabel('像素 Y')
        plt.title('像素坐标系')
        plt.grid(True)
        plt.legend()
        plt.axis('equal')
        
        # 子图2：世界坐标和转换结果
        plt.subplot(1, 2, 2)
        plt.scatter(self.world_points[:, 0], self.world_points[:, 1], 
                   c='blue', label='真实世界坐标', s=50)
        plt.scatter(transformed_points[:, 0], transformed_points[:, 1], 
                   c='green', marker='x', label='转换后的坐标', s=50)
        
        # 添加点编号
        for i, (x, y) in enumerate(self.world_points):
            plt.annotate(f'{i+1}', (x, y), xytext=(5, 5), textcoords='offset points')
        for i, (x, y) in enumerate(transformed_points):
            plt.annotate(f'{i+1}', (x, y), xytext=(-15, 5), textcoords='offset points', color='green')
        
        # 绘制误差线
        for i in range(len(self.world_points)):
            plt.plot([self.world_points[i, 0], transformed_points[i, 0]],
                    [self.world_points[i, 1], transformed_points[i, 1]],
                    'r--', alpha=0.5, linewidth=1)
        
        plt.xlabel('世界 X')
        plt.ylabel('世界 Y')
        plt.title('世界坐标系（电机坐标系）')
        plt.grid(True)
        plt.legend()
        plt.axis('equal')
        
        plt.tight_layout()
        plt.show()
    
    def save_calibration_results(self, output_file="hand_eye_calibration_results.txt"):
        """保存标定结果到文件"""
        if self.H is None:
            print("请先进行标定")
            return
        
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write("# 手眼标定结果\n")
            f.write(f"# 生成时间: {np.datetime64('now')}\n")
            f.write(f"# 标定点数量: {len(self.pixel_points)}\n")
            f.write(f"# 重投影误差: {self.calculate_reprojection_error()[2]:.6f}\n\n")
            
            f.write("单应性矩阵 H:\n")
            for row in self.H:
                f.write("[" + ", ".join(f"{x:.8f}" for x in row) + "]\n")
            
            f.write(f"\n尺度因子 (像素->世界): {self.scale_factor:.8f}\n")
            f.write(f"尺度因子 (世界->像素): {1/self.scale_factor:.8f}\n\n")
            
            f.write("坐标转换验证:\n")
            f.write("序号, 像素X, 像素Y, 真实世界X, 真实世界Y, 转换世界X, 转换世界Y, 误差\n")
            transformed_points = cv2.perspectiveTransform(
                self.pixel_points.reshape(-1, 1, 2), self.H
            ).reshape(-1, 2)
            
            for i in range(len(self.pixel_points)):
                error = np.linalg.norm(transformed_points[i] - self.world_points[i])
                f.write(f"{i+1}, {self.pixel_points[i,0]:.2f}, {self.pixel_points[i,1]:.2f}, "
                       f"{self.world_points[i,0]:.2f}, {self.world_points[i,1]:.2f}, "
                       f"{transformed_points[i,0]:.2f}, {transformed_points[i,1]:.2f}, "
                       f"{error:.6f}\n")
        
        print(f"标定结果已保存到: {output_file}")
    
    def run_calibration(self):
        """运行完整的手眼标定流程"""
        print("开始手眼标定...")
        
        # 1. 估计单应性矩阵
        self.estimate_homography()
        
        # 2. 计算尺度因子
        self.calculate_scale_factor()
        
        print("\n=== 标定完成 ===")
        print("可以使用以下方法进行坐标转换:")
        print("  pixel_to_world(pixel_x, pixel_y) - 像素坐标转世界坐标")
        print("  world_to_pixel(world_x, world_y) - 世界坐标转像素坐标")
        
        # 测试转换
        test_pixel = self.pixel_points[0]
        test_world = self.world_points[0]
        
        converted_world = self.pixel_to_world(test_pixel[0], test_pixel[1])
        converted_pixel = self.world_to_pixel(test_world[0], test_world[1])
        
        print(f"\n测试转换:")
        print(f"像素坐标 {test_pixel} -> 世界坐标 {converted_world}")
        print(f"世界坐标 {test_world} -> 像素坐标 {converted_pixel}")

# 使用示例
if __name__ == "__main__":
    # 创建手眼标定实例
    hand_eye = HandEyeCalibration()
    
    try:
        # 直接从txt文件读取数据
        hand_eye.load_calibration_data_from_file("calibration_data_20250917_180721.txt")
        
        # 运行标定
        hand_eye.run_calibration()
        
        # 可视化结果
        hand_eye.visualize_calibration()
        
        # 保存标定结果
        hand_eye.save_calibration_results()
        
        # 测试一些坐标转换
        print("\n=== 坐标转换测试 ===")
        test_points = [
            (170, 86),   # 已知点1
            (262, 312),  # 已知点2
            (300, 200),  # 中间点
            (400, 100),  # 边界点
        ]
        
        for pixel_x, pixel_y in test_points:
            world_x, world_y = hand_eye.pixel_to_world(pixel_x, pixel_y)
            print(f"像素 ({pixel_x}, {pixel_y}) -> 世界 ({world_x:.2f}, {world_y:.2f})")
        
        # 测试反向转换
        print("\n=== 反向转换测试 ===")
        test_world_points = [
            (68.0, 100.0),  # 已知点1
            (101.0, 18.0),  # 已知点2
            (120.0, 60.0),  # 中间点
        ]
        
        for world_x, world_y in test_world_points:
            pixel_x, pixel_y = hand_eye.world_to_pixel(world_x, world_y)
            print(f"世界 ({world_x}, {world_y}) -> 像素 ({pixel_x:.2f}, {pixel_y:.2f})")
            
    except FileNotFoundError as e:
        print(f"错误: {e}")
        print("请确保 calibration_data_20250917_180721.txt 文件存在")
    except Exception as e:
        print(f"程序执行出错: {e}")